#pragma once

#include <armadillo>
#include "mpi_helper.h"

using namespace arma;

template<typename>
class RAND {
};

template<>
class RAND<mat> {
 public:
  mpi::communicator& world;
  po::variables_map& vm;
  mat W, H;
  int m, n, k;
  double sparsity;
  int n_part, rank;
  vector<int> r,c;

  RAND<mat>(mpi::communicator& world_, po::variables_map& vm_) 
    : world(world_), vm(vm_) 
  {
    REQUIRE_PARAM("m");
    REQUIRE_PARAM("n");
    REQUIRE_PARAM("randk");
    REQUIRE_PARAM("sparse");
    m = vm["m"].as<int>();
    n = vm["n"].as<int>();
    k = vm["randk"].as<int>();
    sparsity = vm["sparse"].as<double>();
    n_part = SIZE(world);
    rank = world.rank();
  }
  
  void randNMF(mat& Arows, mat& Acols) { // multiple processes
    PART<mat>::split_idx(m, n_part, r);
    PART<mat>::split_idx(n, n_part, c);
    srand(vm["s"].as<int>()); // same seed on all processes
    W = randu<mat>(m,k);
    H = randu<mat>(n,k);
    makeSparse(W, sparsity);
    makeSparse(H, sparsity);
    
    Arows = W.rows(r[rank], rank < n_part-1 ? r[rank+1]-1 : m-1) * trans(H);
    Acols = W * trans(H.rows(c[rank], rank < n_part-1 ? c[rank+1]-1 : n-1));
  }

  void randNMF(mat& A) { // single process
    srand(vm["s"].as<int>());
    W = randu<mat>(m,k);
    H = randu<mat>(n,k);
    makeSparse(W, sparsity);
    makeSparse(H, sparsity);
    A = W * trans(H);
  }

  void makeSparse(mat& A, double sparsity) { // make a matrix sparse
    FORE(i, 0, A.n_rows) {
      FORE(j, 0, A.n_cols) {
	if (randu() > sparsity) A(i,j) = 0;
      }
    }
  }
};
